#!/usr/bin/python
# -*- coding: utf-8 -*-
# @Time: 20-7-17 下午4:38
# @Author: HuangChuQin
# @Desc: 查询指定日期A股所有股票的日数据
# @File: stock_hist_date.py
# @Software: PyCharm

import baostock as bs
import pymongo
import common.util_mysql
import common.utils
import sys
import queue
import threading
import time

# 配置文件
profile = 'profile.ini'
# 日志对象
logger = None
# 配置对象
config = None

# 数据库参数
db_param = {}
# 数据库连接
db_conn = None

# mongodb集合
collections = None
# K线指标字典
field_dict = {}


def my_init_conf():
    global profile
    global logger
    global config
    global db_param
    global collections

    # init confFile
    print("init config file:\t" + profile)
    config = common.utils.init_config(profile)
    if config is None:
        print("init config err, exit")
        sys.exit()
    print("init profile success")

    # init log
    log_config_file = config.get("COMMON", "LOG_CONFIG_PATH")
    log_config_name = config.get("COMMON", "LOG_NAME")
    print("log config : ", log_config_file, log_config_name)
    logger = common.utils.init_logger(log_config_file, log_config_name)
    if logger is None:
        print("init logger err, exit")
        sys.exit()
    print("init logger success")

    # init mysql
    db_param['SERVER'] = config.get("MYSQL", "SERVER")
    db_param['USER'] = config.get("MYSQL", "USER")
    db_param['PWD'] = config.get("MYSQL", "PWD")
    db_param['DBNAME'] = config.get("MYSQL", "DBNAME")

    client = pymongo.MongoClient('mongodb://127.0.0.1/')
    print(client.list_database_names())
    collections = client["stock_warehouse"]
    print(collections.list_collection_names())


def db_conn_create():
    global db_conn
    db_conn = common.util_mysql.db_connect(db_param)
    if db_conn is None:
        logger.error('ConnDbMysql err')
        sys.exit()
    logger.info('connect mysql ok')


def db_conn_destroy():
    common.util_mysql.db_disconnect(db_conn)
    logger.info("destroy db_conn ok")


def get_stock_index():
    global field_dict
    (resultSet, rowCount, columnName) = common.util_mysql.fetch_multi_row(
        db_conn, "SELECT index_name FROM stock_base_info.stock_index_info_daily where index_platform='baostock'"
    )
    for result in resultSet:
        (index_name,) = result
        field_dict[index_name] = None


class StockDataThread(threading.Thread):
    def __init__(self, thread_id, name, q, date):
        threading.Thread.__init__(self)
        self.thread_id = thread_id
        self.name = name
        self.q = q
        self.date = date

    def run(self):
        print("开启线程：" + self.name)
        queue_lock.acquire()
        data_dict = get_stock_data(self.name, self.date, ','.join(list(field_dict.keys())))
        write_stock_data("stock_k_data_daily_" + self.date.replace("-", ""), data_dict)
        queue_lock.release()
        print("退出线程：" + self.name)

queue_lock = None
work_queue = None


def get_all_stock(date):
    global queue_lock
    global work_queue

    lg = bs.login()

    print('login respond error_code:' + lg.error_code)
    print('login respond  error_msg:' + lg.error_msg)

    rs = bs.query_all_stock(date)

    print('query_all_stock respond error_code:' + rs.error_code)
    print('query_all_stock respond  error_msg:' + rs.error_msg)

    queue_lock = threading.Lock()
    work_queue = queue.Queue(10)
    threads = []
    thread_id = 1

    # 创建新线程
    for datum in rs.data:
        thread = StockDataThread(thread_id, datum[0], work_queue, date)
        thread.start()
        threads.append(thread)
        thread_id += 1

    # 等待所有线程完成
    for t in threads:
        t.join()
    print("退出主线程")

    bs.logout()


def get_stock_data(code, date, fields):
    rs = bs.query_history_k_data_plus(
        code, fields, date, date, frequency="d"
    )
    print("query_history_k_data_plus, code:%s, error_code:%s, error_msg:%s" % (code, rs.error_code, rs.error_msg))

    data_dict = field_dict.copy()
    index = 0
    for key in data_dict:
        data_dict[key] = rs.data[0][index]
        index += 1
    return data_dict


def write_stock_data(collection_name, data_dict):
    collection = collections[collection_name]
    collection.create_index([("code", 1)], unique=True)
    collection.replace_one({"code": data_dict["code"]}, data_dict, upsert=True)


def main():
    print("######## 开始 "+str(time.time()))
    get_stock_index()
    get_all_stock("2020-07-16")
    print("######## 结束 "+str(time.time()))


if __name__ == "__main__":
    my_init_conf()
    db_conn_create()

    main()

    db_conn_destroy()

    # 退出程序
    logger.info("run end, exit")
